#include "jrm.h"

#include <assert.h>
#include <cuda.h>
#include <stdio.h>
#include <stdlib.h>

#ifndef MAX_THREADS
#ifndef MIN_THREADS
#ifndef WARP_BASED
#ifndef WARP_BASED2
#define WARP_BASED
#endif
#endif
#endif
#endif

#define IS_INITIALIZED  (MA != NULL)
#define WAS_EXECUTED    (VX != NULL)

#define ABS(X)          ((X) > 0 ? (X) : -(X))

#define MIN(X, Y)       ((X) < (Y) ? (X) : (Y))

#define MAX(X, Y)       ((X) > (Y) ? (X) : (Y))

/* Matrices and vectors order */
static int ORDER    = -1;
/* Row to be tested */
static int ROW_TEST = -1;
/* Iterations limit */
static int ITE_MAX  = -1;
/* Iterations for each thread */
static int *ITER    = NULL;

/* Error precision */
static float ERROR  = 0;

/* Matrices and vectors in main memory */
static float *MA    = NULL;
static float *MA_S  = NULL;
static float *VB    = NULL;
static float *VB_S  = NULL;
static float *VX    = NULL;

/* Matrices and vectors in device memory */
static float *D_MA_S   = NULL;
static float *D_VB_S   = NULL;
static float *D_VX     = NULL;

/* Iterations for each thread in device memory */
static int   *D_ITER   = NULL;

/* Device properties */
static cudaDeviceProp dev_props;

void jrm_initialize(char *input)
{
    int i, j;

    if (IS_INITIALIZED)
        jrm_delete();

    /* Parse input file */
    ORDER = (int) strtol(input, &input, 10);
    ROW_TEST = (int) strtol(input, &input, 10);
    ERROR = (float) strtod(input, &input);
    ITE_MAX = (int) strtol(input, &input, 10);

    assert(ORDER > 0);

    /* Get device properties */
    cudaGetDeviceProperties(&dev_props, 0);

    /* Allocate memory for A, A*, b and b* */
    MA = (float *) malloc(ORDER * ORDER * sizeof(float));
    MA_S = (float *) malloc(ORDER * ORDER * sizeof(float));

    VB = (float *) malloc(ORDER * sizeof(float));
    VB_S = (float *) malloc(ORDER * sizeof(float));

    /* Fill matrices and vectors */
    for (i = 0; i < ORDER; i++)
    {
        for (j = 0; j < ORDER; j++)
        {
            MA[i * ORDER + j] = (float) strtod(input, &input);
        }
    }

    for (i = 0; i < ORDER; i++)
        VB[i] = (float) strtod(input, &input);
}

void jrm_delete()
{
    if (!IS_INITIALIZED)
        return;

    /*
     * Free all memory and set original values in variables.
     */

    free(MA);
    free(MA_S);
    free(VB);
    free(VB_S);

    if (WAS_EXECUTED)
    {
        free(VX);
        free(ITER);

        cudaFree(D_VX);
        cudaFree(D_MA_S);
        cudaFree(D_VB_S);
        cudaFree(D_ITER);
    }

    ORDER    = -1;
    ROW_TEST = -1;
    ITE_MAX  = -1;
    ITER     = NULL;

    ERROR  = 0;

    MA = NULL;
    VB = NULL;
    MA_S = NULL;
    VB_S = NULL;
    VX  = NULL;
}

/*
 * Kernel
 */
__global__ void jrm_cuda_execute(float *D_VX, float *D_MA_S,
        float *D_VB_S, int *D_ITER, int MAX_ITER, float ERROR, int ORDER)
{
    __shared__ float numerator, denominator;

    int i, j, iter;
    float x, a, b;

    /* Which x[i] the thread will calculate */
    i = blockIdx.x * blockDim.x + threadIdx.x;

    ERROR = ERROR * ERROR;

    if (i > ORDER)
        return;

    /* Cache main values of b* and A* for this thread */
    b = D_VB_S[i];
    a = D_MA_S[i * ORDER + i];

    for (iter = 0; iter < MAX_ITER; iter++)
    {
        numerator = 0.0f;
        denominator = 0.0f;

        /* Calculate new x value */
        x = b;
        for (j = 0; j < ORDER; j++)
            x -= D_MA_S[i * ORDER + j] * D_VX[j];
        x += a * x;

        /* Update error */
        if (ABS(x - D_VX[i]) > numerator)
            numerator = ABS(x - D_VX[i]);

        if (ABS(x) > denominator)
            denominator = ABS(x);

        /* Update global x */
        D_VX[i] = x;

        /* Wait other threads */
        __syncthreads();

        if (numerator / denominator < ERROR)
            break;
    }

    D_ITER[i] = iter;
}

static void jrm_execute()
{
    int i, j;
    int grid_size, block_size;

    /* Allocate memory for vector x and iterations for each thread */
    VX = (float *) malloc(ORDER * sizeof(float));
    ITER = (int *) malloc(ORDER * sizeof(int));

    /* Allocate device memory */
    cudaMalloc(&D_VX, ORDER * sizeof(float));
    cudaMalloc(&D_MA_S, ORDER * ORDER * sizeof(float));
    cudaMalloc(&D_VB_S, ORDER * sizeof(float));
    cudaMalloc(&D_ITER, ORDER * sizeof(int));

    /* Calculate matrix A* using matrix A */
    for (i = 0; i < ORDER; i++)
        for
            (j = 0; j < ORDER; j++)
            if (i != j)
                MA_S[i * ORDER + j] = MA[i * ORDER + j] / MA[i * ORDER + i];
            else
                MA_S[i * ORDER + j] = 0;

    /* Calculate vector b* using vector b and matrix A */
    for (i = 0; i < ORDER; i++)
        VB_S[i] = VB[i] / MA[i * ORDER + i];

    /* Copy values of x, A* and b* to device */
    cudaMemcpy(D_VX, VB_S, ORDER * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(D_MA_S, MA_S, ORDER * ORDER * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(D_VB_S, VB_S, ORDER * sizeof(float), cudaMemcpyHostToDevice);


    /* Calculate block and grid size depending on chosen approach */
#ifdef MAX_THREADS
    block_size = MIN(ORDER, dev_props.maxThreadsPerBlock);
    grid_size = ORDER / block_size + (ORDER % block_size != 0);
#endif
#ifdef MIN_THREADS
    block_size = MIN(ORDER, dev_props.warpSize);
    grid_size = ORDER / block_size + (ORDER % block_size != 0);
#endif
#ifdef WARP_BASED2
    int warp_size = dev_props.warpSize;

    grid_size = warp_size;
    block_size = ORDER / warp_size + 1;
    block_size = (block_size / warp_size + 1) * warp_size;
#endif
#ifdef WARP_BASED
    int warp_size = dev_props.warpSize;

    block_size = ORDER / warp_size + 1;
    block_size = (block_size / warp_size + 1) * warp_size;
    block_size *= 2;
    grid_size = ORDER / block_size + 1;
#endif

    printf("%d blocks and %d threads per block\n", grid_size, block_size);

    /* Execute cuda kernel */
    jrm_cuda_execute<<<grid_size, block_size>>>(D_VX, D_MA_S, D_VB_S, D_ITER, ITE_MAX, ERROR, ORDER);

    /* Copy calculated values from device to host */
    cudaMemcpy(ITER, D_ITER, ORDER * sizeof(int), cudaMemcpyDeviceToHost);
    cudaMemcpy(VX, D_VX, ORDER * sizeof(float), cudaMemcpyDeviceToHost);
}

void jrm_evaluate(int *iter, int *row_test, float *estimated, float *real)
{
    int i;

    if (!IS_INITIALIZED)
        return;

    if (!WAS_EXECUTED)
        jrm_execute();

    *estimated = 0;
    *iter = 0;

    /* Calculate estimated value in ROW_TEST and the maximum number of iterations */
    for (i = 0; i < ORDER; i++) {
        *estimated += MA[ROW_TEST * ORDER + i] * VX[i];
        if (*iter < ITER[i])
            *iter = ITER[i];
    }

    *real = VB[ROW_TEST];
    *row_test = ROW_TEST;
}
